"""
中继服务器API接口模块
"""

import requests
import json
from typing import Dict, Optional
from config import RELAY_API_BASE, CONNECTION_TIMEOUT

class RelayAPI:
    def __init__(self, base_url: str = RELAY_API_BASE):
        self.base_url = base_url
        self.timeout = CONNECTION_TIMEOUT
    
    def register(self, client_info: str = "", ssh_port: int = 22) -> Dict:
        """
        注册被控端连接
        
        Args:
            client_info: 客户端信息
            ssh_port: SSH端口
            
        Returns:
            包含密钥和隧道端口的字典
        """
        url = f"{self.base_url}/register"
        data = {
            "client_info": client_info,
            "ssh_port": ssh_port
        }
        
        try:
            response = requests.post(url, json=data, timeout=self.timeout)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            return {"status": "error", "message": f"网络错误: {str(e)}"}
    
    def connect(self, key: str) -> Dict:
        """
        验证密钥并获取连接信息
        
        Args:
            key: 8位密钥
            
        Returns:
            包含连接信息的字典
        """
        url = f"{self.base_url}/connect"
        data = {"key": key.upper()}
        
        try:
            response = requests.post(url, json=data, timeout=self.timeout)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            return {"status": "error", "message": f"网络错误: {str(e)}"}
    
    def heartbeat(self, key: str) -> Dict:
        """
        发送心跳保持连接
        
        Args:
            key: 8位密钥
            
        Returns:
            心跳响应
        """
        url = f"{self.base_url}/heartbeat"
        data = {"key": key.upper()}
        
        try:
            response = requests.post(url, json=data, timeout=self.timeout)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            return {"status": "error", "message": f"网络错误: {str(e)}"}
    
    def unregister(self, key: str) -> Dict:
        """
        注销连接
        
        Args:
            key: 8位密钥
            
        Returns:
            注销响应
        """
        url = f"{self.base_url}/unregister"
        data = {"key": key.upper()}
        
        try:
            response = requests.post(url, json=data, timeout=self.timeout)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            return {"status": "error", "message": f"网络错误: {str(e)}"}
    
    def get_status(self) -> Dict:
        """
        获取服务器状态
        
        Returns:
            服务器状态信息
        """
        url = f"{self.base_url}/status"
        
        try:
            response = requests.get(url, timeout=self.timeout)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            return {"status": "error", "message": f"网络错误: {str(e)}"}

